from gensim.models import Word2Vec
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
from mpl_toolkits.mplot3d import Axes3D  # 用于3D绘图

# 设置字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 支持中文字体
plt.rcParams['axes.unicode_minus'] = False

# 加载模型
model = Word2Vec.load("word2vec.model")

# 提取需要可视化的词汇和对应的词向量
words = list(model.wv.index_to_key)[:50]  # 选择前50个常见词汇
word_vectors = model.wv[words]

# 使用PCA将词向量降维到3D
pca = PCA(n_components=3)
word_vectors_3d = pca.fit_transform(word_vectors)

# 创建3D图形
fig = plt.figure(figsize=(12, 10))
ax = fig.add_subplot(111, projection='3d')

# 绘制散点图
ax.scatter(word_vectors_3d[:, 0], word_vectors_3d[:, 1], word_vectors_3d[:, 2], marker='o')

# 添加词汇标签
for i, word in enumerate(words):
    ax.text(word_vectors_3d[i, 0], word_vectors_3d[i, 1], word_vectors_3d[i, 2], word)

# 设置标题和轴标签
ax.set_title("Word2Vec 词向量可视化 (3D PCA)")
ax.set_xlabel("PCA 组件 1")
ax.set_ylabel("PCA 组件 2")
ax.set_zlabel("PCA 组件 3")
plt.show()
